# import paddlex as pdx
import paddle
import paddle.nn as nn
import paddle.nn.functional as F


class Embeder(nn.Layer):
    def __init__(self, word_num, word_size):
        super().__init__()
        self.embedder = nn.Embedding(word_num, word_size)

    def forward(self, input):
        output = self.embedder(input)

        return output


class Encoder(nn.Layer):
    def __init__(self, input_dim, hidden_dim, num_layers,
                 dropout):  # 定义一个初始化函数，在Python类中规定，函数的第一个参数是实例对象本身，并且约定俗成，把其名字写为self。其作用相当于java中的this，表示当前类的对象，可以调用当前类中的属性和方法。
        super().__init__()
        self.word_size = input_dim
        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers, dropout=dropout, time_major=True)

    def forward(self, input, seq_len):  # 该神经网络向前传播函数
        hidden_output, (hidden_h, hidden_c) = self.lstm(input,
                                                        sequence_length=seq_len)  # pytorch的lstm的输入和input（seq_len,batch,input_size）
        # 说人话就是input(每个句子有几个单词（时间步长，循环次数），有几个句子，每个句子的单词由多少维度的向量表示)
        return hidden_output, hidden_h, hidden_c


class Decoder(nn.Layer):

    def __init__(self, input_dim, output_dim, hidden_dim, num_layers,
                 dropout):  # 定义一个初始化函数，在Python类中规定，函数的第一个参数是实例对象本身，并且约定俗成，把其名字写为self。其作用相当于java中的this，表示当前类的对象，可以调用当前类中的属性和方法。
        super().__init__()
        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers, dropout=dropout, time_major=True)
        self.linner = nn.Linear(hidden_dim, output_dim)  # 输出层与输入层的维度相同

    def forward(self, input, hidden_h, hidden_c):  # 该神经网络向前传播函数
        hidden_output, (hidden_h, hidden_c) = self.lstm(input, (hidden_h, hidden_c))  # 这里的变量都是一维的因为我需要一个一个循环
        output = self.linner(hidden_output)  # 这里是在将tensor的输出转为一个常数
        return output, hidden_output, hidden_h, hidden_c


class Attention(nn.Layer):
    def __init__(self,
                 hidden_size):  # 定义一个初始化函数，在Python类中规定，函数的第一个参数是实例对象本身，并且约定俗成，把其名字写为self。其作用相当于java中的this，表示当前类的对象，可以调用当前类中的属性和方法。
        super().__init__()
        self.hidden_size = hidden_size
        self.Wa = nn.Linear(hidden_size * 2, hidden_size)
        self.VT = nn.Linear(hidden_size, 1)

    def forward(self, s, input_output, mask):
        # mask[batch_size,step_time]
        # input_output[step_time,batch_size,hidden_size]
        # h[step_time,batch_size,hidden_size]

        h = input_output.multiply(mask.transpose([1, 0]).unsqueeze(-1).astype("float32"))



        na = self.score(s, h)  # na:[[权值1]，[权值2]...] shape[step_time,batch_size]
        a = F.softmax(na, axis=0)  # na:[概率1，概率2...]shape[step_time,batch_size]
        # a.shape[step_time,batch_size]
        a = a.unsqueeze([-1]) * paddle.ones([self.hidden_size])
        # a.shape[step_time,batch_size,hidden_size]
        r = h.multiply(a)
        # r[step_time,batch_size,hidden_size]
        c = paddle.sum(r, axis=0).unsqueeze(0)
        # c[1,batch_size,hidden_size]
        return c

    def score(self, s, h):
        # 此函数为相似度计算 计算某时刻解码器的隐层向量与编码器所有隐层向量的相似度。并以此为权值
        # 关键函数:输入S与H 从而得到e，进而得出权值，以c为s0，h1为s1 h2为s2....，此为相关性计算
        # 输出：权值
        # h[step_time,batch_size,hidden_size]
        # h=h.transpose([1,0,2])
        # s[1,batch_size,hidden_size]
        # s=s.transpose([1,0,2])
        s = paddle.expand(s, h.shape)
        # s[step_time,batch_size,hidden_size]
        res = self.VT(paddle.tanh(self.Wa(paddle.concat([s, h], axis=-1))))
        # res[step_time,batch_size,1]
        res = res.squeeze(-1)
        return res


class Seq2seq(nn.Layer):
    def __init__(self, hidden_dim, num_layers,
                 dropout):  # 定义一个初始化函数，在Python类中规定，函数的第一个参数是实例对象本身，并且约定俗成，把其名字写为self。其作用相当于java中的this，表示当前类的对象，可以调用当前类中的属性和方法。
        super().__init__()
        self.embeder = Embeder(10255, 64)
        self.encoder = Encoder(64, hidden_dim, num_layers, dropout)
        self.decoder = Decoder(hidden_dim + 64, 10255, hidden_dim, num_layers, dropout)
        self.attention = Attention(hidden_dim)

    def forward(self, inputs, seq_len, mask, labels):  # 该神经网络向前传播函数
        seq_len = seq_len.squeeze(1)
        inputs_embed = self.embeder(inputs)
        maxlen = inputs_embed.shape[1]
        inputs_embed = inputs_embed.transpose([1, 0, 2])
        hidden_output, C_h, C_c = self.encoder.forward(inputs_embed, seq_len)  # 语义向量中h,c 这样分开写便于我愚昧的大脑理解
        sinput = paddle.ones(shape=[inputs.shape[0], 1], dtype='int64') * 3  # batch-size，1
        deinput_lastlabel = paddle.concat([sinput, labels], axis=1, name=None)
        deinput_last = paddle.slice(deinput_lastlabel, axes=[1], starts=[0], ends=[-1])
        deinput_last_embed = self.embeder(deinput_last)
        deinput_last_embed = deinput_last_embed.transpose([1, 0, 2])
        s = self.get_s(C_h)
        for i in range(maxlen):

            c = self.attention(s, hidden_output, mask)  # 得到第一阶段语义注意力语义C

            deinput = paddle.concat([deinput_last_embed[i].unsqueeze(0), c], axis=-1, name=None).astype("float32")  #
            output, de_hidden_output, C_h, C_c = self.decoder.forward(deinput, C_h, C_c)  # 这里的output输出是3维度
            s = de_hidden_output
            if i > 0:
                outputs = paddle.concat([outputs, output], axis=0, name=None)
            else:
                outputs = output
        return outputs.transpose([1, 0, 2])

    def get_s(self, hidden_h):  # 得到S的初始值
        s = hidden_h[-1]  # [2(num_layers), batch_size, hidden_size]
        return s.unsqueeze(0)

    def infer(self, inputs, seq_len, mask):
        inputs = inputs.unsqueeze(0)
        mask = mask.unsqueeze(0)
        inputs_embed = self.embeder(inputs).transpose([1, 0, 2])
        hidden_output, C_h, C_c = self.encoder.forward(inputs_embed, seq_len)  # 语义向量中h,c 这样分开写便于我愚昧的大脑理解
        sinput = paddle.ones(shape=[inputs.shape[0], 1], dtype='int64') * 3
        sinput_embed = self.embeder(sinput).transpose([1, 0, 2])
        s = self.get_s(C_h)
        for i in range(100):
            c = self.attention(s, hidden_output, mask)  # 得到第一阶段语义注意力语义C
            deinput = paddle.concat([sinput_embed, c], axis=-1, name=None).astype("float32")  #
            output, de_hidden_output, C_h, C_c = self.decoder.forward(deinput, C_h, C_c)  # 这里的output输出是3维度
            next_sinput = paddle.argmax(output, axis=-1)
            sinput = paddle.to_tensor(next_sinput, dtype="int64")
            sinput_embed = self.embeder(sinput)
            s = de_hidden_output
            if i > 0:
                outputs = paddle.concat([outputs, output], axis=0, name=None)
            else:
                outputs = output
            if next_sinput[0][0] == 1:
                break
        outsputs=outputs.transpose([1, 0, 2])
        outputs=paddle.argmax(outsputs,-1)
        return outputs
